import torch
import numpy as np
from utils.loss import get_loss
import config.zunYiParameter as para
from sklearn import metrics
from sklearn.metrics import f1_score
from sklearn.metrics import accuracy_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import classification_report


class AvgMeter(object):
    def __init__(self, num=40):
        self.num = num
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
        self.losses = []

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
        self.losses.append(val)

    def show(self):
        return torch.mean(torch.stack(self.losses[np.maximum(len(self.losses)-self.num, 0):]))


class Metirc():
    def __init__(self, real_mask, pred_mask):
        """
        适用于2D和3D
        :param real_mask: 标签
        :param pred_mask: 预测结果
        """
        self.real_mask = real_mask
        self.pred_mask = pred_mask
        self.labels=[i for i in range(para.n_class)]#可能出现的类别标签
        self.__average__="samples"
        self.__pos_label__=1
        self.zero_division=0
    def get_confusion_matrix(self):
        y_true = self.real_mask.reshape(1, -1).squeeze()
        y_pred = self.pred_mask.reshape(1, -1).squeeze()
        confusion_matrix = metrics.confusion_matrix(y_true, y_pred,labels = self.labels)
        return confusion_matrix
    def get_mean_iou(self):
        confusion_matrix = self.get_confusion_matrix()
        intersection = np.diag(confusion_matrix)  # 交集
        union = np.sum(confusion_matrix, axis = 1) + np.sum(confusion_matrix, axis = 0) - np.diag(
            confusion_matrix)  # 并集
        with np.errstate(divide='ignore', invalid='ignore'):#忽略除法中出现0的警告
            IoU = intersection / union  # 交并比，即IoU
        #MIoU = np.mean(IoU)  # 计算MIoU
        MIoU = np.nanmean(IoU)
        return IoU,MIoU

    def precision(self):
        precision = precision_score(self.real_mask, self.pred_mask, labels = self.labels,zero_division = self.zero_division,
                                    average = self.__average__, pos_label = self.__pos_label__)
        return precision

    def recall(self):
        recall = recall_score(self.real_mask, self.pred_mask, labels = self.labels,zero_division = self.zero_division
                              ,average = self.__average__, pos_label = self.__pos_label__)
        return recall

    def accuracy(self):
        OA = accuracy_score(self.real_mask, self.pred_mask)
        # OA=accuracy_score(y_true, y_pred, normalize=False)
        return OA

    def F1(self):
        f1 = f1_score(self.real_mask, self.pred_mask , labels = self.labels,zero_division = self.zero_division,
                     average = self.__average__)
        return f1
    def classification_report(self):
        result=classification_report(y_true = self.real_mask,y_pred = self.pred_mask,target_names = self.labels,labels = self.labels,output_dict = True,zero_division = self.zero_division)
        return result
def get_confusion_matrix(real_mask,pred_mask,n_class=2):
    '''
   适用于部分图像没有某个类别标签的情况。例如2D肿瘤分割，很多图片没有肿瘤。
   这时候要先计算所有图像的混淆矩阵，再相加，然后利用总的混淆矩阵计算各项评估指标
   '''
    y_true = real_mask.reshape(1, -1).squeeze()
    y_pred = pred_mask.reshape(1, -1).squeeze()
    confusion_matrix = metrics.confusion_matrix(y_true, y_pred,labels = [0,1])
    return confusion_matrix

class CaculateLastTypeFromConfusionMatrix():
    '''
    适用于部分图像没有某个类别标签的情况。例如2D肿瘤分割，很多图片没有肿瘤。
    这时候要先计算所有图像的混淆矩阵，再相加，然后利用总的混淆矩阵计算各项评估指标
    '''
    def __init__(self,confusion_matrix):
        '''
        [[2.51330000e+04,1.64190000e+04 ],
        [6.76200000e+03, 1.54092358e+08]]
        '''
        self.tn, self.fp, self.fn, self.tp = confusion_matrix.ravel()  # tn, fp, fn, tp
        # self.tp = confusion_matrix[0][0]
        # self.fn = confusion_matrix[0][1]
        # self.fp = confusion_matrix[1][0]
        # self.tn = confusion_matrix[1][1]
    def get_accuracy(self):
        return (self.tp+self.tn)/(self.tp+self.fn+self.fp+self.tn)
    def get_iou(self):
        return self.tp/(self.fn+self.fp+self.tp)
    def get_dice(self):
        return 2*self.tp/(self.fp+2*self.tp+self.fn)
    def get_precision(self):
        return self.tp/(self.tp+self.fp)
    def get_recall(self):
        return self.tp/(self.tp+self.fn)
    def get_f1(self):
        return 2*self.tp/(2*self.tp+self.fp+self.fn)
# def mean_iou_np(y_true, y_pred, **kwargs):
#     """
#     compute mean iou for binary segmentation map via numpy
#     只计算了类别1的iou
#     """
#     y_true = y_true.flatten()
#     y_pred = y_pred.flatten()
#     intersection = np.sum(np.abs(y_pred * y_true))
#     mask_sum = np.sum(np.abs(y_true)) + np.sum(np.abs(y_pred))
#     union = mask_sum - intersection
#
#     smooth = .0001
#     iou = (intersection + smooth) / (union + smooth)
#
#     return iou
#
#
# def mean_dice_np(y_true, y_pred, **kwargs):
#     """
#     compute mean dice for binary segmentation map via numpy
#     """
#     y_true=y_true.flatten()
#     y_pred=y_pred.flatten()
#     intersection = np.sum(np.abs(y_pred * y_true))
#     mask_sum = np.sum(np.abs(y_true)) + np.sum(np.abs(y_pred))
#     smooth = .0001
#     dice = 2 * (intersection + smooth) / (mask_sum + smooth)
#     return dice

def evaluate(model,test_loader):
    lossf=get_loss()
    model.eval()
    mean_loss = []
    test_loader = test_loader
    loss_bank = []
    n_class=2
    confusion_matrix_total = np.zeros((n_class, n_class))
    index=0
    for image, gt in test_loader:
        image = image.cuda()

        with torch.no_grad():
            _, _, res = model(image)
        # print(image.shape,res.shape,torch.tensor(gt).unsqueeze(0).unsqueeze(0).shape)
        # loss = structure_loss(res, torch.tensor(gt).unsqueeze(0).unsqueeze(0).cuda())
        loss = lossf(res, gt.cuda())
        res = res.sigmoid().data.cpu().numpy().squeeze()
        gt = 1*(gt>0.5)
        res = 1*(res > 0.5)

        gt=gt.cpu().numpy()

        confusion_matrix = get_confusion_matrix(gt,res)
        confusion_matrix_total += confusion_matrix
        loss_bank.append(loss.item())
        # if 1 in gt:
        #     dice = mean_dice_np(gt, res)
        #     iou = mean_iou_np(gt, res)
        #     acc = np.sum(res == gt) / (res.shape[0]*res.shape[1])
        #
        #     dice_bank.append(dice)
        #     iou_bank.append(iou)
        #     acc_bank.append(acc)
    print(confusion_matrix_total)
    Caculate = CaculateLastTypeFromConfusionMatrix(confusion_matrix_total)

    iou = Caculate.get_iou()
    dice = Caculate.get_dice()
    accuracy=Caculate.get_accuracy()
    precision=Caculate.get_precision()
    recall=Caculate.get_recall()
    f1=Caculate.get_f1()
    mean_loss=np.mean(loss_bank)
    # print('test Loss: {:.4f}, Dice: {:.4f}, IoU: {:.4f}, Acc: {:.4f}'.
    #     format( mean_loss, dice, iou, acc))

    result_dict={'iou':iou,'dice':dice,'accuracy':accuracy,
                 'precision':precision,'recall':recall,'f1':f1,
                 'mean loss:':mean_loss,}
    return result_dict

if __name__ == '__main__':
    a=[0,1,1,0,1,1,1,0]
    b=[1,1,1,0,0,0,0,1]
    # self.tp = confusion_matrix[0][0]
    # self.fn = confusion_matrix[0][1]
    # self.fp = confusion_matrix[1][0]
    # # self.tn = confusion_matrix[1][1]
    result = [[2.51330000e+04,1.64190000e+04 ],
              [6.76200000e+03, 1.54092358e+08]]
    # result=metrics.confusion_matrix(a,b,labels = [0,1])
    #self.tn, self.fp, self.fn, self.tp = confusion_matrix.ravel()  # tn, fp, fn, tp
    print(result)
    Caculate = CaculateLastTypeFromConfusionMatrix(np.array(result))
    iou = Caculate.get_iou()
    dice = Caculate.get_dice()
    accuracy = Caculate.get_accuracy()
    precision = Caculate.get_precision()
    recall = Caculate.get_recall()
    f1 = Caculate.get_f1()
    # print('test Loss: {:.4f}, Dice: {:.4f}, IoU: {:.4f}, Acc: {:.4f}'.
    #     format( mean_loss, dice, iou, acc))
    result_dict = {'iou': iou, 'dice': dice, 'accuracy': accuracy,
                   'precision': precision, 'recall': recall, 'f1': f1,
                   }
    print(dice,f1)